from keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D, concatenate
from keras.optimizers import Adam
from keras.models import Model
from keras.losses import binary_crossentropy
from keras.callbacks import ModelCheckpoint, TensorBoard
import cv2
import glob
import numpy as np
import tensorflow as tf

# 定义 YOLOv5 的基础模块
def conv_block(input_tensor, num_filters):
    x = Conv2D(num_filters, (3, 3), padding='same', activation='relu')(input_tensor)
    x = Conv2D(num_filters, (3, 3), padding='same', activation='relu')(x)
    return x

# 下采样模块
def downsample_block(input_tensor, num_filters):
    x = conv_block(input_tensor, num_filters)
    p = MaxPooling2D((2, 2))(x)
    return x, p

# 上采样模块
def upsample_block(input_tensor, concat_tensor, num_filters):
    x = UpSampling2D((2, 2))(input_tensor)
    x = concatenate([x, concat_tensor], axis=-1)
    x = conv_block(x, num_filters)
    return x

# YOLOv5 模型
def yolov5_model(input_shape=(416, 416, 3), num_classes=80):
    # 输入层
    inputs = Input(shape=input_shape)

    # 编码器
    x1, p1 = downsample_block(inputs, 64)
    x2, p2 = downsample_block(p1, 128)
    x3, p3 = downsample_block(p2, 256)
    x4, p4 = downsample_block(p3, 512)
    x5, p5 = downsample_block(p4, 1024)

    # 中间层
    x5 = conv_block(x5, 1024)

    # 解码器
    u1 = upsample_block(x5, x4, 512)
    u2 = upsample_block(u1, x3, 256)
    u3 = upsample_block(u2, x2, 128)
    u4 = upsample_block(u3, x1, 64)

    # 输出层
    outputs = Conv2D(num_classes, (1, 1), activation='sigmoid')(u4)

    model = Model(inputs=inputs, outputs=outputs)
    return model

import keras.backend as K
def yolo_loss(y_true, y_pred):
    # 1. 坐标损失
    coord_mask = y_true[..., 4:5]
    coord_loss = coord_mask * (K.square(y_pred[..., :4] - y_true[..., :4]))

    # 2. 置信度损失
    conf_mask = y_true[..., 4:5]
    conf_loss_true = conf_mask * K.square(1 - y_pred[..., 4:5])
    conf_loss_false = (1 - conf_mask) * K.square(0 - y_pred[..., 4:5])
    conf_loss = conf_loss_true + conf_loss_false

    # 3. 类别损失
    class_mask = y_true[..., 5:]
    class_loss = class_mask * K.square(y_pred[..., 5:] - y_true[..., 5:])

    # 汇总损失
    coord_loss_sum = K.sum(coord_loss)
    conf_loss_sum = K.sum(conf_loss)
    class_loss_sum = K.sum(class_loss)
    total_loss = coord_loss_sum + conf_loss_sum + class_loss_sum
    return total_loss

def yolo_accuracy(y_true, y_pred):
    # 计算预测正确的边界框数量
    correct_boxes = tf.reduce_sum(tf.cast(tf.equal(tf.argmax(y_true[..., 5:], axis=-1), tf.argmax(y_pred[..., 5:], axis=-1)), tf.float32))
    # 总边界框数量
    total_boxes = tf.cast(tf.reduce_prod(tf.shape(y_true)[0:3]), tf.float32)
    # 准确率
    return correct_boxes / total_boxes

from keras.preprocessing.image import ImageDataGenerator
def generator(image_directory,target_size,batch_size,dict_data,seed):
    # 定义图像生成器
    datagen_image = ImageDataGenerator(**dict_data)
    datagen_label = ImageDataGenerator(**dict_data)
    data = zip(
        datagen_image.flow_from_directory(
            image_directory,
            classes=['image'],
            target_size=target_size,
            batch_size=batch_size,       #样本数不易过大
            seed=seed),
        datagen_label.flow_from_directory(
            image_directory,
            classes=["label"],
            color_mode='grayscale',
            target_size=target_size,
            batch_size=batch_size,       #样本数不易过大
            seed=seed)
        )
    for (img,label) in data:
        # cv2.imshow("image", img[0][0,:,:,:])
        # cv2.imshow("label", label[0][0,:,:,:])
        # cv2.waitKey(0)
        yield img[0],label[0]
        
if __name__ == "__main__":
    path = 'deep_net\\test'
    img = cv2.imread((glob.glob(path+"\\image\\*.jpg")+glob.glob(path+"\\image\\*.png"))[0])
    img = cv2.resize(img, (img.shape[1]//16*16, img.shape[0]//16*16))
    model = yolov5_model()
    if True:
        model.summary()
        data = generator(path,img.shape[:2], 2, dict(
            rescale=1./255,rotation_range=10,width_shift_range=0.2,height_shift_range=0.2,
            shear_range=0.2,zoom_range=0.2,horizontal_flip=True,fill_mode='constant'),1)
        model.compile(optimizer=Adam(learning_rate=0.01), loss=yolo_loss, metrics=[yolo_accuracy])
        model.fit(data, steps_per_epoch=15, epochs=100, callbacks=
                [
                    TensorBoard('deep_net/logs',write_images=True,histogram_freq=1,write_graph=False),
                    ModelCheckpoint("deep_net\\param\\unet.keras"),
                ])
    else:
        model.load_weights("deep_net\\param\\unet.keras")
        for file in glob.glob(path+"\\image\\*.jpg")+glob.glob(path+"\\image\\*.png"):
            new_img = cv2.imread(file)/255.
            new_img = cv2.resize(new_img,(img.shape[1],img.shape[0]))          
            new_img = new_img.reshape(1,*img.shape[:2],3)
            new_img = model.predict(new_img)
            cv2.imshow("",new_img[0,:,:,0])
            cv2.waitKey(1)